// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: jshin

#pragma once

#include <cstdlib>
#include <memory>
#include <string>
#include <thread>

#include <boost/asio/bind_executor.hpp>
#include <boost/asio/io_context_strand.hpp>
#include <boost/asio/io_service.hpp>
#include <boost/asio/post.hpp>

#include "coin2/app/executor.pb.h"
#include "coin2/exchange/base/executor/async_executor.h"
#include "coin2/exchange/base/feed/counter/counter.h"
#include "coin2/flow/topic_string.h"
#include "naviasync/ipc/coin_reactor.h"

namespace coin2::exchange::base::executor {

using coin2::exchange::base::feed::counter::FeedCounter;

class LiveTopicInjector {
 public:
  LiveTopicInjector(
      boost::asio::io_context* main_ioc,
      SharedTopicRecordSubscriber* subscriber)
      : main_ioc_(main_ioc),
        subscriber_(subscriber) {}
  LiveTopicInjector() {}

  virtual void push(std::shared_ptr<TopicRecordPacket> packet) {
    boost::asio::post(
        *main_ioc_,
        std::bind(
            &SharedTopicRecordSubscriber::OnTopicRecord,
            subscriber_,
            packet));
  }

  virtual void RegisterIoc(const std::set<std::string>& topics) {
  }

  boost::asio::io_context* main_ioc_ = nullptr;
  SharedTopicRecordSubscriber* subscriber_ = nullptr;
};

class LiveTopicCounterInjector : public LiveTopicInjector {
 public:
  LiveTopicCounterInjector(
      FeedCounter* counter,
      std::unique_ptr<LiveTopicInjector> inject)
      : counter_(counter),
        inject_(std::move(inject)) {
  }

  virtual void push(std::shared_ptr<TopicRecordPacket> packet) {
    inject_->push(packet);
    counter_->ReceiveOneMsg();
  }

  virtual void RegisterIoc(const std::set<std::string>& topics) {
    inject_->RegisterIoc(topics);
  }

 private:
  FeedCounter* counter_;
  std::unique_ptr<LiveTopicInjector> inject_;
};

class LiveTopicParallelInjector : public LiveTopicInjector {
 public:
  LiveTopicParallelInjector(
      boost::asio::io_context* run_ioc,
      SharedTopicRecordSubscriber* subscriber)
      : run_ioc_(run_ioc),
        subscriber_(subscriber) {}

  virtual void push(std::shared_ptr<TopicRecordPacket> packet) {
    auto t = coin2::flow::TopicStringConfig::StripWokerFromConnectionTopic(packet->topic_name);
    auto iter = topic2strand_.find(t);
    CHECK(iter != topic2strand_.end());
    boost::asio::post(
        *(iter->second),
        std::bind(
            &SharedTopicRecordSubscriber::OnTopicRecord,
            subscriber_,
            packet));
  }

  virtual void RegisterIoc(const std::set<std::string>& topics) override {
    for (const auto& topic : topics) {
      auto t = coin2::flow::TopicStringConfig::StripWokerFromConnectionTopic(topic);
      auto iter = topic2strand_.find(t);
      if (iter == topic2strand_.end()) {
        auto strand = std::make_unique<boost::asio::io_context::strand>(*run_ioc_);
        topic2strand_[t] = strand.get();
        strands_.emplace_back(std::move(strand));
      }
    }
  }

 private:
  boost::asio::io_context* run_ioc_ = nullptr;
  SharedTopicRecordSubscriber* subscriber_ = nullptr;
  std::map<std::string, boost::asio::io_context::strand*> topic2strand_;
  std::vector<std::unique_ptr<boost::asio::io_context::strand>> strands_;
};

class CoinReactor : public ::naviasync::ipc::CoinReactor<naviasync::ipc::AsioReactorSubscribeMode::kIoCtxOnly> {
 public:
  using Base =
      ::naviasync::ipc::CoinReactor<naviasync::ipc::AsioReactorSubscribeMode::kIoCtxOnly>;
  using Base::Base;
};

class LiveAsyncExecutor : public AsyncExecutor, public SharedTopicRecordSubscriber {
 public:
  explicit LiveAsyncExecutor(const ::coin2::app::LiveExecutorConfig& config)
      : config_(config),
        main_ioc_(1),
        raw_ioc_(config.raw_threads()),
        stats_ioc_(1),
#ifdef DISABLE_REACTOR_INTEGRATION
        writer_ioc_(nullptr) {
#else
        writer_ioc_(nullptr),
        main_reactor_(CoinReactor::Context{}, main_ioc_) {
#endif
    CHECK_GE(config_.raw_threads(), 1);
    if (config.writer_threads() > 0) {
      writer_ioc_ = std::make_unique<boost::asio::io_context>(config.writer_threads());
    }
    if (config.parser_threads() > 0) {
      parser_ioc_ = std::make_unique<boost::asio::io_context>(config.parser_threads());
      topic_injector_ = std::make_unique<LiveTopicParallelInjector>(parser_ioc_.get(), this);
    } else {
      topic_injector_ = std::make_unique<LiveTopicInjector>(&main_ioc_, this);
    }
  }

  void Start() override;
  void Stop() override;

  boost::asio::io_context* MainIoc() { return &main_ioc_; }
  boost::asio::io_context* RawIoc() { return &raw_ioc_; }
  boost::asio::io_context* WriterIoc() { return writer_ioc_.get(); }
  boost::asio::io_context* StatsIoc() { return &stats_ioc_; }
  boost::asio::io_context* ParserIoc() { return parser_ioc_.get(); }
  LiveTopicInjector& injector() { return *topic_injector_; }
  CoinReactor& main_reactor() { return main_reactor_; }
  void EnableCounter(FeedCounter* counter) {
    counter_ = counter;
    topic_injector_.reset(new LiveTopicCounterInjector(counter, std::move(topic_injector_)));
  }

  void OnTopicRecord(std::shared_ptr<TopicRecordPacket> packet) override;

 protected:
  const ::coin2::app::LiveExecutorConfig config_;
  boost::asio::io_context main_ioc_;
  boost::asio::io_context raw_ioc_;
  boost::asio::io_context stats_ioc_;
  std::unique_ptr<boost::asio::io_context> writer_ioc_;
  std::unique_ptr<boost::asio::io_context> parser_ioc_;
  std::unique_ptr<LiveTopicInjector> topic_injector_;
  FeedCounter* counter_ = nullptr;
  CoinReactor main_reactor_;
};
}  // namespace coin2::exchange::base::executor
